import geopandas as gpd
import pandas as pd
import trimesh
import numpy as np
import os
import time
from tqdm import tqdm
from shapely.ops import unary_union
from shapely.validation import make_valid
from shapely import is_valid
from shapely.errors import TopologicalError, GEOSException

def optimized_infringement_farmland_to_3d(
        geojson_path,
        output_path,
        scale=1000000,
        export_format='glb',
        extrusion_height=1.0,
        simplify_tolerance_ratio=0.00001
):
    """
    侵权耕地GeoJSON转Unity 3D模型（修复参数错误+子物体拆分版）
    核心修复：移除trimesh导出不支持的compression参数，确保GLB/OBJ导出正常
    子物体功能：每个地块/子区生成独立网格，Unity导入后自动识别为独立子物体
    """
    # ---------------------- 1. 基础初始化与输入校验 ----------------------
    start_total_time = time.time()
    print(f"【开始】侵权耕地3D模型生成，时间：{time.strftime('%Y-%m-%d %H:%M:%S')}")

    # 路径校验与目录创建
    if not os.path.exists(geojson_path):
        print(f"【错误】输入GeoJSON不存在：{geojson_path}")
        return
    output_dir = os.path.dirname(output_path)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"【提示】创建输出目录：{output_dir}")

    # 导出格式预处理（仅保留trimesh完全支持的格式）
    supported_formats = ['glb', 'gltf', 'obj', 'stl', 'ply']
    export_format = export_format.lower() if export_format.lower() in supported_formats else 'glb'
    # 自动补全文件后缀（避免格式不匹配）
    if not output_path.lower().endswith(f'.{export_format}'):
        output_path = f"{output_path}.{export_format}"

    # 碰撞检测库预加载（可选功能，不影响导出）
    try:
        import fcl
        fcl_available = True
        print("【提示】启用精确碰撞检测（python-fcl）")
    except ImportError:
        fcl_available = False
        print("【提示】启用轻量化边界框碰撞检测")

    # ---------------------- 2. 数据读取与字段预处理 ----------------------
    print("\n【步骤1/5】读取GeoJSON并处理字段...")
    try:
        gdf = gpd.read_file(geojson_path, crs="EPSG:4326", ignore_geometry=False)
        total_raw = len(gdf)
        print(f"【成功】读取{total_raw}个地块，字段：{list(gdf.columns)}")
    except Exception as e:
        print(f"【错误】读取GeoJSON失败：{str(e)}")
        return

    # 核心字段校验（适配侵权耕地字段）
    core_fields = ["LIN_CHANGM", "LIN_BAN2", "XIAO_BAN2", "OBJECTID"]
    missing_fields = [f for f in core_fields if f not in gdf.columns]
    if missing_fields:
        print(f"【错误】缺少核心字段：{missing_fields}，无法生成唯一名称")
        return

    # 字段批量处理（解决空值、类型问题）
    gdf[core_fields] = gdf[core_fields].fillna("未知")  # 空值填充
    gdf["OBJECTID"] = gdf["OBJECTID"].astype(str).str.strip()  # OBJECTID转字符串（避免下标报错）
    # 林班/小班保留前置零（如0247不转为247，保证名称一致性）
    for field in ["LIN_BAN2", "XIAO_BAN2"]:
        if pd.api.types.is_numeric_dtype(gdf[field]):
            max_len = gdf[field].astype(str).str.len().max()
            gdf[field] = gdf[field].astype(str).str.zfill(int(max_len))
        else:
            gdf[field] = gdf[field].astype(str).str.strip()
    # 面积/地类字段补全（防止原始数据缺失）
    if "面积" not in gdf.columns:
        gdf["面积"] = gdf.geometry.area.round(4)
    else:
        gdf["面积"] = pd.to_numeric(gdf["面积"], errors="coerce").fillna(0.0).round(4)
    if "dlmc" not in gdf.columns:
        gdf["dlmc"] = "侵权耕地"
    else:
        gdf["dlmc"] = gdf["dlmc"].astype(str).str.strip()

    # 示例验证（确保字段处理正确）
    sample = gdf.iloc[0]
    print(f"【示例】地块名：{sample['LIN_CHANGM']}{sample['LIN_BAN2']}{sample['XIAO_BAN2']}_OBJ{sample['OBJECTID'][:6]}")

    # ---------------------- 3. 批量几何修复 ----------------------
    print(f"\n【步骤2/5】修复几何错误（共{len(gdf)}个地块）...")
    valid_geoms = []
    for _, row in tqdm(gdf.iterrows(), total=len(gdf), desc="几何修复进度"):
        geom = row['geometry']
        if geom is None or geom.is_empty:
            valid_geoms.append(None)
            continue

        try:
            # 修复无效几何（如自相交、孔洞）
            if not is_valid(geom):
                fixed_geom = make_valid(geom)
                # MultiPolygon保留最大子区（避免生成微小碎片网格）
                if fixed_geom.geom_type == "MultiPolygon":
                    fixed_geom = sorted(fixed_geom.geoms, key=lambda p: p.area if p.exterior else 0, reverse=True)[0]
                valid_geoms.append(fixed_geom)
            else:
                valid_geoms.append(geom)
        except Exception as e:
            valid_geoms.append(None)

    # 过滤无效地块（仅保留可生成网格的有效数据）
    gdf['geometry'] = valid_geoms
    gdf_valid = gdf[gdf.geometry.notna() & ~gdf.geometry.is_empty].copy()
    print(f"【完成】几何修复：{len(gdf)} → {len(gdf_valid)}个有效地块")
    if len(gdf_valid) == 0:
        print("【错误】无有效地块，终止生成")
        return

    # ---------------------- 4. 坐标批量计算 ----------------------
    print(f"\n【步骤3/5】计算坐标偏移与缩放...")
    all_coords = []
    for geom in gdf_valid['geometry']:
        # 统一处理Polygon/MultiPolygon（已保留最大子区）
        poly = geom if geom.geom_type == "Polygon" else geom.geoms[0]
        if poly.exterior is not None:
            all_coords.extend([coord[:2] for coord in poly.exterior.coords])

    # 计算偏移原点（避免Unity坐标过大，影响显示）
    all_coords_np = np.array(all_coords, dtype=np.float64)
    min_lon, min_lat = all_coords_np.min(axis=0)
    offset = np.array([min_lon, min_lat], dtype=np.float64)
    simplify_tolerance = scale * simplify_tolerance_ratio  # 动态简化阈值

    print(f"【坐标配置】偏移（{min_lon:.6f}, {min_lat:.6f}）| 缩放：{scale} | 简化阈值：{simplify_tolerance:.4f}")

    # ---------------------- 5. 生成独立网格（核心：Unity子物体） ----------------------
    print(f"\n【步骤4/5】生成独立3D网格（共{len(gdf_valid)}个地块）...")
    scene = trimesh.Scene()
    valid_meshes = []
    existing_meshes = []

    # 三角化引擎配置（提升网格生成速度）
    try:
        import mapbox_earcut
        os.environ["TRIMESH_TRIANGULATE_ENGINE"] = "earcut"
        print("【提示】使用高效三角化引擎：mapbox_earcut")
    except ImportError:
        os.environ["TRIMESH_TRIANGULATE_ENGINE"] = "default"
        print("【提示】使用默认三角化引擎")

    # 逐地块生成独立网格（每个网格对应Unity一个子物体）
    for _, (_, row) in tqdm(enumerate(gdf_valid.iterrows()), total=len(gdf_valid), desc="网格生成进度"):
        # 1. 生成唯一地块名（避免Unity子物体重名）
        base_name = f"{row['LIN_CHANGM'].strip()}{row['LIN_BAN2'].strip()}{row['XIAO_BAN2'].strip()}_OBJ{row['OBJECTID'][:6]}"

        # 2. 处理几何（支持MultiPolygon子区，每个子区一个网格）
        geom = row['geometry']
        polygons = list(geom.geoms) if geom.geom_type == "MultiPolygon" else [geom]

        # 3. 逐子区生成网格
        for sub_idx, poly in enumerate(polygons):
            mesh_name = f"{base_name}_SUB{sub_idx}"  # 子区名称（如xxx_SUB0）
            if poly.exterior is None:
                continue  # 跳过无外边界的无效子区

            # 4. 坐标转换（经纬度→平面坐标，适配Unity）
            coords_2d = np.array([coord[:2] for coord in poly.exterior.coords], dtype=np.float64)
            if len(coords_2d) < 4:  # 跳过顶点不足的多边形（至少4个顶点构成闭合图形）
                continue
            planar_coords = (coords_2d - offset) * scale  # 偏移+缩放

            # 5. 几何简化（减少网格面数，提升Unity性能）
            path = trimesh.load_path(planar_coords)
            if len(planar_coords) > 200:  # 仅复杂多边形执行简化
                path = path.simplify(tolerance=simplify_tolerance)
            if len(path.vertices) < 3:  # 简化后顶点不足，跳过
                continue

            # 6. 生成3D网格（挤压高度，处理异常情况）
            try:
                # 优先使用trimesh自带挤压（效率高）
                mesh = path.extrude(height=extrusion_height, fix_normals=False)
                # 处理挤压后返回列表的情况（部分几何会生成多个网格）
                if isinstance(mesh, list):
                    mesh = mesh[0] if mesh else None
            except Exception as e:
                # 备选方案：手动构建完整3D网格（避免挤压失败）
                bottom_v = np.column_stack([planar_coords, np.zeros(len(planar_coords))])  # 底面顶点
                top_v = np.column_stack([planar_coords, np.full(len(planar_coords), extrusion_height)])  # 顶面顶点
                all_v = np.vstack([bottom_v, top_v])  # 合并顶底顶点

                # 构建底面三角面
                bottom_f = np.array([[0, i, i+1] for i in range(1, len(bottom_v)-1)], dtype=np.int32)
                # 构建顶面三角面（反向顺序避免法线错误）
                top_f = np.array([
                    [len(bottom_v), len(bottom_v)+i+1, len(bottom_v)+i]
                    for i in range(1, len(top_v)-1)
                ], dtype=np.int32)
                # 构建侧面三角面
                side_f = []
                for i in range(len(bottom_v)-1):
                    v0, v1, v2, v3 = i, i+1, len(bottom_v)+i+1, len(bottom_v)+i
                    side_f.extend([[v0, v1, v2], [v0, v2, v3]])
                side_f = np.array(side_f, dtype=np.int32)

                # 合并所有面，过滤无效面
                all_f = np.vstack([bottom_f, top_f, side_f]) if (
                        len(bottom_f) > 0 and len(top_f) > 0
                ) else None
                if all_f is None or len(all_f) < 4:  # 至少4个面（正方体最小面数）
                    continue
                mesh = trimesh.Trimesh(vertices=all_v, faces=all_f, process=False)

            # 过滤无效网格（确保Unity能正常导入）
            if mesh is None or len(mesh.vertices) < 6 or len(mesh.faces) < 4:
                continue

            # 7. 轻量化碰撞检测（避免网格重叠）
            if existing_meshes and len(existing_meshes) % 5 == 0:
                if fcl_available:
                    if check_collision_fcl(mesh, existing_meshes[-5:]):
                        mesh.apply_translation([1e-5, 1e-5, 0])  # 轻微偏移
                else:
                    current_bbox = mesh.bounds
                    overlap = False
                    for exist_mesh in existing_meshes[-3:]:
                        exist_bbox = exist_mesh.bounds
                        if (current_bbox[0][0] < exist_bbox[1][0] and
                                current_bbox[1][0] > exist_bbox[0][0] and
                                current_bbox[0][1] < exist_bbox[1][1] and
                                current_bbox[1][1] > exist_bbox[0][1]):
                            overlap = True
                            break
                    if overlap:
                        mesh.apply_translation([1e-5, 1e-5, 0])

            # 8. 添加元数据（Unity中可读取地块属性）
            mesh.name = mesh_name
            mesh.metadata = {
                "地类名称": row['dlmc'],
                "地块面积": row['面积'],
                "林场名称": row['LIN_CHANGM'],
                "林班号": row['LIN_BAN2'],
                "小班号": row['XIAO_BAN2'],
                "OBJECTID": row['OBJECTID']
            }

            # 9. 加入场景（关键：每个网格单独添加，确保Unity子物体独立）
            valid_meshes.append(mesh)
            existing_meshes.append(mesh)
            scene.add_geometry(mesh)

    print(f"【完成】网格生成：{len(gdf_valid)}个地块 → {len(valid_meshes)}个独立网格")
    if not valid_meshes:
        print("【错误】无有效网格，终止导出")
        return

    # ---------------------- 6. 最终导出（彻底移除不兼容参数） ----------------------
    print(f"\n【步骤5/5】导出{export_format}模型...")
    try:
        # 核心修复：仅保留trimesh所有版本都支持的基础参数（无compression/include_metadata）
        scene.export(
            file_obj=output_path,
            file_type=export_format  # 仅指定格式，不添加任何多余参数
        )

        # 生成统计信息
        total_time = round(time.time() - start_total_time, 2)
        speed = round(len(valid_meshes) / total_time, 2) if total_time > 0 else len(valid_meshes)
        print(f"\n【🎉 生成完成】===== 结果汇总 =====")
        print(f"文件路径：{output_path}")
        print(f"统计：{len(valid_meshes)}个独立子物体 | 耗时：{total_time}秒 | 速度：{speed}个/秒")
        print(f"Unity提示：导入后展开根节点，每个网格是独立子物体，可单独操作")
    except Exception as e:
        print(f"【错误】{export_format}导出失败：{str(e)}")
        # 备选导出OBJ（兼容性最强，trimesh无参数冲突）
        fallback_path = output_path.rsplit(".", 1)[0] + ".obj"
        try:
            scene.export(fallback_path, file_type="obj")  # OBJ导出同样不添加多余参数
            print(f"【备选成功】已导出OBJ格式：{fallback_path}")
        except Exception as fe:
            print(f"【错误】备选导出失败：{str(fe)}")


def check_collision_fcl(mesh, existing_meshes):
    """精确碰撞检测（仅检查最近网格，减少计算）"""
    import fcl
    # 构建当前网格BVH
    current_bvh = fcl.BVHModel()
    current_bvh.beginModel(len(mesh.vertices), len(mesh.faces))
    current_bvh.addSubModel(mesh.vertices, mesh.faces)
    current_bvh.endModel()
    current_obj = fcl.CollisionObject(current_bvh)

    # 检查最近网格
    for exist_mesh in existing_meshes:
        exist_bvh = fcl.BVHModel()
        exist_bvh.beginModel(len(exist_mesh.vertices), len(exist_mesh.faces))
        exist_bvh.addSubModel(exist_mesh.vertices, exist_mesh.faces)
        exist_bvh.endModel()
        exist_obj = fcl.CollisionObject(exist_bvh)

        # 简化碰撞请求
        req = fcl.CollisionRequest(num_max_contacts=1, enable_contact=False)
        res = fcl.CollisionResult()
        fcl.collide(current_obj, exist_obj, req, res)
        if res.is_collision:
            return True
    return False


# ---------------------- 执行入口（仅需修改路径） ----------------------
if __name__ == "__main__":
    # 【配置路径】替换为你的侵权耕地GeoJSON和输出路径
    INPUT_GEOJSON = r"C:\Users\ma\Desktop\山河屯数据\耕地区域\侵权耕地.geojson"
    OUTPUT_MODEL = r"C:\Users\ma\Desktop\山河屯数据\耕地区域\侵权耕地_3D模型.glb"

    # 调用函数（参数可按需调整）
    optimized_infringement_farmland_to_3d(
        geojson_path=INPUT_GEOJSON,
        output_path=OUTPUT_MODEL,
        scale=1000000,          # 适配Unity坐标，无需修改
        export_format='glb',    # 推荐GLB（Unity最优），可选obj
        extrusion_height=1.2,   # 模型高度（根据场景调整）
        simplify_tolerance_ratio=0.00001  # 简化阈值（越小越精细）
    )
